# -*- coding: utf-8 -*-
"""
超跌选股筛选器（灵活版）
筛选规则：6个技术指标中满足4个即可 + 流动性过滤

技术指标（6选4）：
1. CCI < -100: 商品通道指数超卖
2. KDJ < 30: 随机指标超卖
3. WR > 80: 威廉指标超卖
4. RSI < 30: 相对强弱指数超卖
5. 收盘价 < 布林下轨: 价格偏离正常区间
6. BIAS < -8%: 10日乖离率超跌

流动性过滤（避免庄股）：
- 排除ST、*ST股票
- 成交量 > 近20日平均成交量的30%
- 换手率 > 0.5% (暂时跳过，因数据源不提供)

使用方法：
    python oversold_screener_flexible.py              # 筛选所有股票
    python oversold_screener_flexible.py --limit 100  # 限制筛选数量（测试用）
    python oversold_screener_flexible.py --min-conditions 4  # 最少满足4个条件（默认）
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
import argparse

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'oversold_screener_flexible.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class OversoldScreenerFlexible:
    """超跌选股筛选器（灵活版）"""
    
    def __init__(self, min_conditions: int = 4):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)
        self.min_conditions = min_conditions  # 最少满足几个条件
    
    def get_stock_codes(self, limit: int = None, specific_stock: str = None) -> list:
        """获取所有活跃股票代码（排除ST股票）"""
        try:
            if specific_stock:
                sql = """
                SELECT stock_code, stock_name 
                FROM new_stock_info 
                WHERE stock_code = %s AND status = 'active'
                  AND stock_name NOT LIKE '%ST%'
                  AND stock_name NOT LIKE '%退%'
                """
                results = self.db.execute_query(sql, (specific_stock,))
            else:
                sql = """
                SELECT stock_code, stock_name 
                FROM new_stock_info 
                WHERE status = 'active'
                  AND stock_name NOT LIKE '%ST%'
                  AND stock_name NOT LIKE '%退%'
                ORDER BY stock_code
                """
                if limit:
                    sql += f" LIMIT {limit}"
                results = self.db.execute_query(sql)
            
            return [(row['stock_code'], row['stock_name']) for row in results]
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_daily_data(self, stock_code: str, days: int = 50) -> pd.DataFrame:
        """获取股票最近N天的日线数据"""
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, 
               volume, amount, change_pct, turnover_rate
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换数值列
            for col in ['open_price', 'high_price', 'low_price', 'close_price', 
                       'volume', 'amount', 'change_pct', 'turnover_rate']:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 按日期升序排列
            df = df.sort_values('trade_date').reset_index(drop=True)
            
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 日线数据失败: {e}")
            return pd.DataFrame()
    
    @staticmethod
    def calculate_cci(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算CCI指标"""
        if df.empty or len(df) < period:
            df['cci'] = np.nan
            return df
        
        tp = (df['high_price'] + df['low_price'] + df['close_price']) / 3
        ma = tp.rolling(window=period).mean()
        md = tp.rolling(window=period).apply(lambda x: np.abs(x - x.mean()).mean())
        df['cci'] = (tp - ma) / (0.015 * md)
        
        return df
    
    @staticmethod
    def calculate_wr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算WR指标"""
        if df.empty or len(df) < period:
            df['wr'] = np.nan
            return df
        
        highest_high = df['high_price'].rolling(window=period).max()
        lowest_low = df['low_price'].rolling(window=period).min()
        df['wr'] = (highest_high - df['close_price']) / (highest_high - lowest_low + 0.0001) * 100
        
        return df
    
    @staticmethod
    def calculate_rsi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算RSI指标"""
        if df.empty or len(df) < period:
            df['rsi'] = np.nan
            return df
        
        delta = df['close_price'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        avg_loss = avg_loss.replace(0, 0.0001)
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - (100 / (1 + rs))
        
        return df
    
    @staticmethod
    def calculate_kdj(df: pd.DataFrame, n: int = 9, m1: int = 3, m2: int = 3) -> pd.DataFrame:
        """计算KDJ指标"""
        if df.empty or len(df) < n:
            df['kdj_k'] = np.nan
            df['kdj_d'] = np.nan
            df['kdj_j'] = np.nan
            return df
        
        low_min = df['low_price'].rolling(window=n).min()
        high_max = df['high_price'].rolling(window=n).max()
        rsv = 100 * (df['close_price'] - low_min) / (high_max - low_min + 0.0001)
        df['kdj_k'] = rsv.ewm(span=m1, adjust=False).mean()
        df['kdj_d'] = df['kdj_k'].ewm(span=m2, adjust=False).mean()
        df['kdj_j'] = 3 * df['kdj_k'] - 2 * df['kdj_d']
        
        return df
    
    @staticmethod
    def calculate_bollinger(df: pd.DataFrame, period: int = 20, std_dev: float = 2.0) -> pd.DataFrame:
        """计算布林带"""
        if df.empty or len(df) < period:
            df['boll_mid'] = np.nan
            df['boll_upper'] = np.nan
            df['boll_lower'] = np.nan
            return df
        
        df['boll_mid'] = df['close_price'].rolling(window=period).mean()
        std = df['close_price'].rolling(window=period).std()
        df['boll_upper'] = df['boll_mid'] + std_dev * std
        df['boll_lower'] = df['boll_mid'] - std_dev * std
        
        return df
    
    @staticmethod
    def calculate_bias(df: pd.DataFrame, period: int = 10) -> pd.DataFrame:
        """计算乖离率"""
        if df.empty or len(df) < period:
            df['bias'] = np.nan
            return df
        
        ma = df['close_price'].rolling(window=period).mean()
        df['bias'] = (df['close_price'] - ma) / ma * 100
        
        return df
    
    @staticmethod
    def calculate_volume_avg(df: pd.DataFrame, period: int = 20) -> pd.DataFrame:
        """计算平均成交量"""
        if df.empty or len(df) < period:
            df['volume_avg'] = np.nan
            return df
        
        df['volume_avg'] = df['volume'].rolling(window=period).mean()
        
        return df
    
    def check_liquidity(self, df: pd.DataFrame) -> dict:
        """
        检查流动性条件（避免庄股）
        
        条件：
        1. 成交量 > 近20日平均成交量的30%
        2. 换手率 > 0.5% (如果数据可用)
        
        注意：由于Tushare的daily接口不提供换手率，暂时只检查成交量
        """
        if df.empty or len(df) < 2:
            return {'passed': False, 'reason': '数据不足'}
        
        latest = df.iloc[-1]
        
        # 检查成交量
        if pd.isna(latest['volume_avg']) or latest['volume_avg'] == 0:
            return {'passed': False, 'reason': '成交量数据缺失'}
        
        volume_ratio = latest['volume'] / latest['volume_avg']
        turnover = latest['turnover_rate']
        
        # 流动性过滤 - 成交量检查
        if volume_ratio < 0.3:
            return {'passed': False, 'reason': f'成交量过低(仅为20日均量的{volume_ratio*100:.1f}%)'}
        
        # 换手率检查 - 只有在数据可用时才检查
        # 由于数据库中大部分股票换手率为NULL，暂时跳过此检查
        # if not pd.isna(turnover) and turnover < 0.5:
        #     return {'passed': False, 'reason': f'换手率过低({turnover:.2f}%)'}
        
        return {
            'passed': True,
            'volume_ratio': volume_ratio,
            'turnover_rate': turnover if not pd.isna(turnover) else 0.0
        }
    
    def check_oversold_conditions(self, df: pd.DataFrame) -> dict:
        """
        检查超跌条件（6选4）
        
        6个技术指标：
        1. CCI < -100
        2. KDJ < 30
        3. WR > 80
        4. RSI < 30
        5. 收盘价 < 布林下轨
        6. BIAS < -8%
        """
        if df.empty or len(df) < 2:
            return {
                'is_oversold': False,
                'conditions_met': 0,
                'required_conditions': self.min_conditions,
                'reason': '数据不足'
            }
        
        latest = df.iloc[-1]
        
        # 检查各个条件
        conditions = {
            'CCI': (latest.get('cci', np.nan) < -100, latest.get('cci', np.nan), '<-100'),
            'KDJ': (latest.get('kdj_k', np.nan) < 30, latest.get('kdj_k', np.nan), '<30'),
            'WR': (latest.get('wr', np.nan) > 80, latest.get('wr', np.nan), '>80'),
            'RSI': (latest.get('rsi', np.nan) < 30, latest.get('rsi', np.nan), '<30'),
            'BOLL': (latest.get('close_price', np.nan) < latest.get('boll_lower', np.nan), 
                     f"{latest.get('close_price', 0):.2f}<{latest.get('boll_lower', 0):.2f}", '跌破下轨'),
            'BIAS': (latest.get('bias', np.nan) < -8, latest.get('bias', np.nan), '<-8%')
        }
        
        # 统计满足的条件
        met_conditions = []
        condition_details = []
        
        for name, (met, value, threshold) in conditions.items():
            if met and not pd.isna(value):
                met_conditions.append(name)
                if name == 'BOLL':
                    condition_details.append(f"{name}({value})")
                else:
                    condition_details.append(f"{name}={value:.2f}{threshold}")
        
        conditions_met_count = len(met_conditions)
        is_oversold = conditions_met_count >= self.min_conditions
        
        return {
            'is_oversold': is_oversold,
            'conditions_met': conditions_met_count,
            'required_conditions': self.min_conditions,
            'met_list': met_conditions,
            'condition_details': condition_details,
            'reason': ', '.join(condition_details) if condition_details else '不满足超跌条件',
            'latest_date': latest['trade_date'],
            'close_price': latest['close_price'],
            'cci': latest.get('cci', np.nan),
            'kdj': latest.get('kdj_k', np.nan),
            'wr': latest.get('wr', np.nan),
            'rsi': latest.get('rsi', np.nan),
            'boll_lower': latest.get('boll_lower', np.nan),
            'bias': latest.get('bias', np.nan),
            'volume_ratio': latest.get('volume', 0) / latest.get('volume_avg', 1) if latest.get('volume_avg', 0) > 0 else 0,
            'turnover_rate': latest.get('turnover_rate', 0)
        }
    
    def screen_single_stock(self, stock_code: str, stock_name: str) -> dict:
        """筛选单只股票"""
        try:
            # 获取日线数据
            df = self.get_stock_daily_data(stock_code, days=50)
            
            if df.empty:
                return {
                    'stock_code': stock_code,
                    'stock_name': stock_name,
                    'is_oversold': False,
                    'reason': '无日线数据'
                }
            
            # 计算技术指标
            df = self.calculate_cci(df, period=14)
            df = self.calculate_kdj(df, n=9, m1=3, m2=3)
            df = self.calculate_wr(df, period=14)
            df = self.calculate_rsi(df, period=14)
            df = self.calculate_bollinger(df, period=20, std_dev=2.0)
            df = self.calculate_bias(df, period=10)
            df = self.calculate_volume_avg(df, period=20)
            
            # 检查流动性
            liquidity_check = self.check_liquidity(df)
            if not liquidity_check['passed']:
                return {
                    'stock_code': stock_code,
                    'stock_name': stock_name,
                    'is_oversold': False,
                    'reason': f"流动性不足: {liquidity_check['reason']}"
                }
            
            # 检查超跌条件
            result = self.check_oversold_conditions(df)
            result['stock_code'] = stock_code
            result['stock_name'] = stock_name
            
            return result
            
        except Exception as e:
            logger.error(f"筛选股票 {stock_code} 失败: {e}")
            return {
                'stock_code': stock_code,
                'stock_name': stock_name,
                'is_oversold': False,
                'reason': f'筛选失败: {e}'
            }
    
    def screen_stocks(self, limit: int = None, specific_stock: str = None) -> list:
        """批量筛选股票"""
        stock_list = self.get_stock_codes(limit=limit, specific_stock=specific_stock)
        
        if not stock_list:
            logger.warning("没有找到符合条件的股票")
            return []
        
        logger.info(f"开始筛选 {len(stock_list)} 只股票（已排除ST股票）...")
        
        results = []
        oversold_stocks = []
        
        for i, (stock_code, stock_name) in enumerate(stock_list, 1):
            if i % 100 == 0 or i == len(stock_list):
                logger.info(f"进度: {i}/{len(stock_list)}")
            
            result = self.screen_single_stock(stock_code, stock_name)
            
            if result['is_oversold']:
                oversold_stocks.append(result)
                logger.info(f"✓ 发现超跌股票: {stock_code} {stock_name} - "
                          f"满足{result['conditions_met']}/{result['required_conditions']}个条件 - "
                          f"{', '.join(result['met_list'])}")
            
            results.append(result)
        
        logger.info(f"筛选完成！共发现 {len(oversold_stocks)} 只超跌股票")
        
        return results
    
    def save_results(self, results: list, output_format: str = 'both'):
        """保存筛选结果"""
        oversold_stocks = [r for r in results if r['is_oversold']]
        
        if not oversold_stocks:
            logger.info("没有发现超跌股票，不生成输出文件")
            return
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 输出逗号分隔的股票代码
        stock_codes = [s['stock_code'] for s in oversold_stocks]
        codes_string = ','.join(stock_codes)
        
        print("\n" + "=" * 80)
        print("超跌股票代码（逗号分隔）：")
        print("=" * 80)
        print(codes_string)
        print("=" * 80)
        
        # 保存CSV文件
        if output_format in ['csv', 'both']:
            df = pd.DataFrame(oversold_stocks)
            csv_file = f'oversold_stocks_flexible_{timestamp}.csv'
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"CSV文件已保存: {csv_file}")
        
        # 保存TXT文件
        if output_format in ['txt', 'both']:
            txt_file = f'stock_codes_oversold_flexible_{timestamp}.txt'
            with open(txt_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(stock_codes))
            logger.info(f"股票代码文件已保存: {txt_file}")
        
        # 打印详细信息
        print("\n超跌股票详细信息：")
        print("-" * 160)
        print(f"{'代码':<8} {'名称':<10} {'日期':<12} {'满足':<8} {'条件列表':<35} "
              f"{'CCI':<8} {'KDJ':<8} {'WR':<8} {'RSI':<8} {'BIAS%':<8} {'换手%':<8}")
        print("-" * 160)
        
        for stock in oversold_stocks:
            met_list = ','.join(stock['met_list'])
            print(f"{stock['stock_code']:<8} {stock['stock_name']:<10} "
                  f"{str(stock['latest_date'])[:10]:<12} "
                  f"{stock['conditions_met']}/{stock['required_conditions']:<6} "
                  f"{met_list:<35} "
                  f"{stock['cci']:<8.2f} {stock['kdj']:<8.2f} {stock['wr']:<8.2f} "
                  f"{stock['rsi']:<8.2f} {stock['bias']:<8.2f} {stock['turnover_rate']:<8.2f}")
        
        print("-" * 160)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='超跌选股筛选器（灵活版）')
    parser.add_argument('--limit', type=int, help='限制筛选股票数量（用于测试）')
    parser.add_argument('--stock', type=str, help='指定筛选某只股票代码')
    parser.add_argument('--min-conditions', type=int, default=4,
                       help='最少满足几个条件（默认4，范围3-6）')
    parser.add_argument('--format', type=str, default='both', 
                       choices=['csv', 'txt', 'both'],
                       help='输出格式：csv, txt, both (默认both)')
    
    args = parser.parse_args()
    
    # 验证条件数量
    if args.min_conditions < 3 or args.min_conditions > 6:
        logger.error("最少条件数量必须在3-6之间")
        return 1
    
    try:
        logger.info("=" * 80)
        logger.info("超跌选股筛选器（灵活版）启动")
        logger.info(f"筛选条件: 6个技术指标中满足{args.min_conditions}个即可")
        logger.info("流动性过滤: 排除ST股 + 成交量>20日均30%")
        logger.info("注意: 换手率检查已临时禁用（数据源不提供）")
        logger.info("=" * 80)
        
        # 创建筛选器
        screener = OversoldScreenerFlexible(min_conditions=args.min_conditions)
        
        # 执行筛选
        results = screener.screen_stocks(limit=args.limit, specific_stock=args.stock)
        
        # 保存结果
        screener.save_results(results, output_format=args.format)
        
        logger.info("=" * 80)
        logger.info("筛选完成")
        logger.info("=" * 80)
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}", exc_info=True)
        return 1
    
    return 0


if __name__ == "__main__":
    sys.exit(main())

